import requests
import json

# Ollama API基础URL
oollama_url = "http://localhost:11434"

print("=== 简单Ollama API测试 ===")

# 1. 测试基本连接
try:
    print("\n1. 测试基本连接...")
    response = requests.get(f"{oollama_url}", timeout=5)
    print(f"连接状态码: {response.status_code}")
    print(f"响应内容: {response.text[:100]}...")
except requests.exceptions.ConnectionError:
    print("连接失败: 无法连接到Ollama服务")
except Exception as e:
    print(f"连接异常: {str(e)}")

# 2. 测试列出模型接口
try:
    print("\n2. 测试列出模型接口...")
    response = requests.get(f"{oollama_url}/api/tags", timeout=5)
    if response.status_code == 200:
        models = response.json()
        print(f"可用模型数量: {len(models.get('models', []))}")
        for model in models.get('models', []):
            print(f"- {model.get('name')} (尺寸: {model.get('size')}, 格式: {model.get('format')})")
    else:
        print(f"请求失败: 状态码 {response.status_code}")
        print(f"响应内容: {response.text}")
except Exception as e:
    print(f"获取模型列表异常: {str(e)}")

# 3. 测试聊天接口
try:
    print("\n3. 测试聊天接口...")
    chat_data = {
        "model": "qwen2.5:7b",
        "messages": [
            {"role": "user", "content": "你好，这是一个测试消息。"}
        ],
        "stream": False
    }
    
    response = requests.post(f"{oollama_url}/api/chat", json=chat_data, timeout=10)
    
    if response.status_code == 200:
        chat_response = response.json()
        print(f"聊天响应成功!")
        print(f"响应内容: {chat_response.get('message', {}).get('content')}")
    else:
        print(f"聊天请求失败: 状态码 {response.status_code}")
        print(f"响应内容: {response.text}")
except Exception as e:
    print(f"聊天请求异常: {str(e)}")

# 4. 测试生成接口
try:
    print("\n4. 测试生成接口...")
    generate_data = {
        "model": "qwen2.5:7b",
        "prompt": "写一句简短的问候语",
        "stream": False
    }
    
    response = requests.post(f"{oollama_url}/api/generate", json=generate_data, timeout=10)
    
    if response.status_code == 200:
        generate_response = response.json()
        print(f"生成响应成功!")
        print(f"响应内容: {generate_response.get('response')}")
    else:
        print(f"生成请求失败: 状态码 {response.status_code}")
        print(f"响应内容: {response.text}")
except Exception as e:
    print(f"生成请求异常: {str(e)}")

print("\n=== 测试完成 ===")